import argparse
from osgeo.gdal import Open, GetDriverByName, GDT_Int16
import numpy as np
import scipy.cluster.vq as vq


def isodata(data):
    # TODO isodata cluster
    print("cluster with isodata")
    return data[:, :, 0]


class Classification:
    def __init__(self, args, **kwargs):
        self.input_path = None
        self.output_path = None
        self.method = "kmeans"
        self.num_classes = 5
        self.threshold = 1e-5
        self.kmeans_init = "random"
        if args is not None:
            for key in args.__dict__:
                self[key] = args.__dict__[key]
        for key in kwargs.keys():
            self[key] = kwargs[key]

    def __getitem__(self, item):
        return self.__dict__[item]

    def __setitem__(self, key, value):
        self.__dict__[key] = value

    def __str__(self):
        return f"input={self.input_path},output={self.output_path},method={self.method}"

    def run(self):
        # 读取数据
        img_src = Open(self.input_path)
        proj = img_src.GetProjection()
        trans = img_src.GetGeoTransform()
        n_bands, n_rows, n_cols = img_src.RasterCount, img_src.RasterYSize, img_src.RasterXSize
        n_pixels = n_rows * n_cols
        img = img_src.ReadAsArray()
        # gdal读取的数据是按照【通道，行，列】的方式排列的，将其转换为【行，列，通道】
        img = np.transpose(img, (1, 2, 0))
        # 将数据重新组织为【行*列，通道】的格式，数据的每一行代表一个像素
        data = np.reshape(img, [n_pixels, n_bands])
        # 进行分类
        if self.method == 'isodata':
            # 使用isodata方法进行聚类
            img_label = isodata(data)
        else:
            # 默认使用kmeans方法
            data = data.astype(dtype=np.float32)
            centroids, labels = vq.kmeans2(data=data, k=self.num_classes, thresh=self.threshold, minit=self.kmeans_init,
                                           missing='warn', check_finite=True)
            img_label = np.reshape(labels, (n_rows, n_cols))
        # 输出分类结果
        driver = GetDriverByName("GTiff")
        img_out = driver.Create(self.output_path, n_cols, n_rows, 1, GDT_Int16)
        img_out.SetGeoTransform(trans)
        img_out.SetProjection(proj)
        img_out.GetRasterBand(1).WriteArray(img_label)
        del img_out


if __name__ == '__main__':
    # 接收控制台传入参数
    parser = argparse.ArgumentParser(description='classification parameters')
    parser.add_argument('input_path', type=str, help='input file path')
    parser.add_argument('output_path', type=str, help='output file path')
    parser.add_argument('--method', '-m', type=str, default='kmeans', choices=['kmeans', 'isodata'],
                        help="cluster method, and default is 'kmeans'")
    parser.add_argument('--num_classes', '-nc', type=int, default=5, help="number of classes")
    parser.add_argument('--threshold', '-t', type=float, default=1e-5,
                        help="threshold to terminates the k-means algorithm")
    parser.add_argument('--kmeans_init', '-ki', type=str, default='random',
                        choices=['random', 'points', '++', 'matrix'],
                        help="initialization method for kmeans centroids")
    args = parser.parse_args()
    # print(args.__dict__)
    # 执行分类
    cls = Classification(args=args)
    # print(cls)
    cls.run()
